from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from langchain_core.tools import tool
from langchain_tavily import TavilySearch
from typing import List, Dict, Any, Annotated, TypedDict, Optional
import requests
import base64
import json
from datetime import datetime, timezone, timedelta

from config import settings


# 定义状态类型
class AgentState(TypedDict):
    """代理状态定义"""
    messages: Annotated[List[Dict[str, str]], "聊天消息历史"]
    current_input: Annotated[str, "当前用户输入"]
    message_type: Annotated[str, "消息类型: text, image, file"]
    file_data: Annotated[bytes, "文件数据"]
    response: Annotated[str, "AI响应"]
    ocr_raw_content: Annotated[str, "OCR识别的原始内容"]


class OpenAIChatModel:
    """OpenAI兼容模型封装 - 使用LangChain 1.0通用大模型调用方式，集成联网搜索功能"""

    def __init__(self):
        # 使用LangChain的ChatOpenAI类，统一模型调用接口
        self.client = ChatOpenAI(
            base_url=settings.openai_api_base,
            api_key=settings.openai_api_key,
            model=settings.openai_model,
            temperature=0.7,
            max_tokens=2048,
            timeout=60.0,  # 添加超时设置，60秒
            max_retries=2   # 添加重试机制
        )

        # 初始化Tavily搜索工具
        self.search_tool = None
        self._init_search_tool()

        # 检查API配置是否有效
        self._validate_config()

    def _init_search_tool(self):
        """初始化搜索工具"""
        # 检查是否启用了搜索功能
        if not settings.tavily_search_enabled:
            print("Tavily搜索功能已禁用")
            self.search_tool = None
            return

        # 检查API密钥是否有效配置
        if not settings.tavily_api_key:
            print("Tavily API密钥未配置，搜索功能将不可用")
            self.search_tool = None
            return

        try:
            # 获取北京时间
            beijing_time = self._get_beijing_time()

            # 创建搜索工具，使用新版本的TavilySearch
            self.search_tool = TavilySearch(
                tavily_api_key=settings.tavily_api_key,
                search_depth="advanced",
                max_results=settings.tavily_max_results,
                include_answer=True,
                include_raw_content=False,
                include_images=False
            )
            print(f"Tavily搜索工具初始化成功 - 北京时间: {beijing_time}")
        except Exception as e:
            print(f"Tavily搜索工具初始化失败: {e}")
            self.search_tool = None

    def _get_beijing_time(self) -> str:
        """获取北京时间"""
        # 北京时间是UTC+8
        beijing_tz = timezone(timedelta(hours=8))
        beijing_now = datetime.now(beijing_tz)
        return beijing_now.strftime("%Y-%m-%d %H:%M:%S")

    async def _perform_search(self, query: str) -> Optional[str]:
        """执行联网搜索"""
        if not self.search_tool or not settings.tavily_search_enabled:
            return None

        try:
            # 获取北京时间
            beijing_time = self._get_beijing_time()

            # 添加时间上下文到搜索查询
            search_query = f"{query} 当前北京时间: {beijing_time}"

            # 执行搜索
            search_results = await self.search_tool.ainvoke({"query": search_query})
            if search_results:
                # 格式化搜索结果
                formatted_results = f"\n\n=== 联网搜索结果（北京时间 {beijing_time}） ===\n"

                # 处理Tavily搜索返回的字典结构
                if isinstance(search_results, dict):
                    # 提取主要信息
                    search_query = search_results.get('query', '无查询')
                    answer = search_results.get('answer', '无答案')
                    follow_up_questions = search_results.get(
                        'follow_up_questions')
                    results = search_results.get('results', [])

                    # 添加查询和答案
                    formatted_results += f"查询: {search_query}\n"
                    formatted_results += f"答案: {answer}\n\n"

                    # 添加后续问题（如果有）
                    if follow_up_questions:
                        formatted_results += f"后续问题: {follow_up_questions}\n\n"

                    # 添加详细搜索结果
                    if results and isinstance(results, list):
                        formatted_results += "详细搜索结果:\n"
                        # 限制前N条结果
                        for i, result in enumerate(results[:settings.tavily_max_results], 1):
                            if isinstance(result, dict):
                                title = result.get('title', '无标题')
                                content = result.get('content', '无内容')
                                url = result.get('url', '无链接')

                                formatted_results += f"{i}. {title}\n"
                                formatted_results += f"   内容: {content[:500]}...\n"
                                formatted_results += f"   链接: {url}\n\n"
                            else:
                                # 如果是其他类型，直接转换为字符串
                                formatted_results += f"{i}. {str(result)}\n\n"
                else:
                    # 如果不是字典结构，直接转换为字符串
                    formatted_results += str(search_results)

                formatted_results += "=== 搜索结束 ===\n"
                return formatted_results

        except Exception as e:
            pass

        return None

    def _validate_config(self):
        """验证API配置"""
        if not settings.openai_api_key:
            raise ValueError("通义千问 API密钥未配置")

        if not settings.openai_api_base:
            raise ValueError("通义千问 API基础URL未配置")

    async def chat(self, messages: List[Dict[str, str]], enable_search: bool = True) -> str:
        """聊天对话，支持联网搜索"""
        try:
            # 获取最新的用户消息
            user_message = None
            for msg in reversed(messages):
                if msg["role"] == "user":
                    user_message = msg["content"]
                    break

            # 如果启用搜索且有用户消息，执行联网搜索
            search_results = None
            if enable_search and user_message and settings.tavily_search_enabled:
                search_results = await self._perform_search(user_message)

            # 转换消息格式
            langchain_messages = []
            for msg in messages:
                if msg["role"] == "user":
                    langchain_messages.append(
                        HumanMessage(content=msg["content"]))
                elif msg["role"] == "assistant":
                    langchain_messages.append(
                        AIMessage(content=msg["content"]))
                elif msg["role"] == "system":
                    langchain_messages.append(
                        SystemMessage(content=msg["content"]))

            # 如果有搜索结果，添加到系统提示中
            if search_results:
                # 创建包含搜索结果的系统提示
                search_prompt = f"""
                当前对话已启用联网搜索功能。以下是基于用户问题的实时搜索结果：
                {search_results}

                请基于这些搜索结果和对话上下文，为用户提供准确、及时的回答。
                """
                # 在消息列表开头添加系统提示
                langchain_messages.insert(
                    0, SystemMessage(content=search_prompt))

            response = await self.client.ainvoke(langchain_messages)

            # 检查响应是否有效
            if not response or not hasattr(response, 'content'):
                return "抱歉，AI服务暂时不可用，请稍后再试"

            return response.content
        except Exception as e:
            # 提供更友好的错误信息
            error_msg = str(e)
            if "choices" in error_msg and "null" in error_msg:
                return "AI服务配置有误：API端点可能不支持OpenAI格式，请检查API配置。当前配置：{} - {}".format(
                    settings.openai_api_base, settings.openai_model
                )
            elif "401" in error_msg or "unauthorized" in error_msg.lower():
                return "API密钥无效或已过期，请检查配置"
            elif "404" in error_msg or "not found" in error_msg.lower():
                return "API端点或模型不存在，请检查配置"
            elif "timeout" in error_msg.lower():
                return "AI服务请求超时，请检查网络连接"
            else:
                return f"AI服务暂时不可用: {error_msg}"

    async def chat_stream(self, messages: List[Dict[str, str]], enable_search: bool = True):
        """流式聊天对话 - 返回生成器，支持联网搜索"""
        try:
            # 获取最新的用户消息
            user_message = None
            for msg in reversed(messages):
                if msg["role"] == "user":
                    user_message = msg["content"]
                    break

            # 如果启用搜索且有用户消息，执行联网搜索
            search_results = None
            if enable_search and user_message and settings.tavily_search_enabled:
                search_results = await self._perform_search(user_message)

            # 转换消息格式
            langchain_messages = []
            for msg in messages:
                if msg["role"] == "user":
                    langchain_messages.append(
                        HumanMessage(content=msg["content"]))
                elif msg["role"] == "assistant":
                    langchain_messages.append(
                        AIMessage(content=msg["content"]))
                elif msg["role"] == "system":
                    langchain_messages.append(
                        SystemMessage(content=msg["content"]))

            # 如果有搜索结果，添加到系统提示中
            if search_results:
                # 创建包含搜索结果的系统提示
                search_prompt = f"""
                当前对话已启用联网搜索功能。以下是基于用户问题的实时搜索结果：
                {search_results}

                请基于这些搜索结果和对话上下文，为用户提供准确、及时的回答。
                """
                # 在消息列表开头添加系统提示
                langchain_messages.insert(
                    0, SystemMessage(content=search_prompt))

            # 使用流式调用 - 注意：astream返回的是异步生成器，不需要await
            response_stream = self.client.astream(langchain_messages)

            async for chunk in response_stream:
                if hasattr(chunk, 'content') and chunk.content:
                    yield chunk.content

        except Exception as e:
            # 提供更友好的错误信息
            error_msg = str(e)
            if "choices" in error_msg and "null" in error_msg:
                yield "AI服务配置有误：API端点可能不支持OpenAI格式，请检查API配置。当前配置：{} - {}".format(
                    settings.openai_api_base, settings.openai_model
                )
            elif "401" in error_msg or "unauthorized" in error_msg.lower():
                yield "API密钥无效或已过期，请检查配置"
            elif "404" in error_msg or "not found" in error_msg.lower():
                yield "AI服务请求超时，请检查网络连接"
            else:
                yield f"AI服务暂时不可用: {error_msg}"


class GeneralOCR:
    """通用OCR模型封装 - 使用LangChain 1.0通用大模型调用方式"""

    def __init__(self):
        # 先检查API配置是否有效
        self._validate_config()

        # 使用LangChain的ChatOpenAI类，统一模型调用接口
        self.client = ChatOpenAI(
            base_url=settings.ocr_api_base,
            api_key=settings.ocr_api_key,
            model=settings.ocr_model,
            temperature=0.1,  # OCR任务需要较低的温度
            max_tokens=1000,
            timeout=30.0,  # 添加超时设置，30秒
            max_retries=2   # 添加重试机制
        )

    def _validate_config(self):
        """验证API配置"""
        if not settings.ocr_api_key:
            raise ValueError("OCR_API_KEY is required")

        if not settings.ocr_api_base:
            raise ValueError("OCR_API_BASE is required")

    async def recognize_text(self, image_data: bytes, user_prompt: str = None) -> str:
        """识别图片中的文字 - 使用LangChain多模态消息格式

        Args:
            image_data: 图片数据字节
            user_prompt: 用户输入的文字提示，如果为None则使用默认提示
        """
        try:
            # 验证图像数据
            if not image_data or len(image_data) == 0:
                return "图像数据为空，请重新上传有效的图片文件"

            # 检查图像数据格式
            try:
                from PIL import Image
                import io
                # 尝试打开图像验证格式
                image = Image.open(io.BytesIO(image_data))

                # 保存原始格式信息
                original_format = image.format

                # 转换为RGB格式以确保兼容性（PNG等格式可能包含透明度）
                if image.mode != 'RGB':
                    image = image.convert('RGB')

                # 根据原始格式选择合适的输出格式
                if original_format and original_format.upper() in ['PNG', 'GIF']:
                    # 对于PNG和GIF格式，保持PNG格式以保留质量
                    output = io.BytesIO()
                    image.save(output, format='PNG', optimize=True)
                    image_data = output.getvalue()
                    image_format = 'png'
                else:
                    # 对于其他格式，使用JPEG
                    output = io.BytesIO()
                    image.save(output, format='JPEG', quality=85)
                    image_data = output.getvalue()
                    image_format = 'jpeg'

            except Exception as img_error:
                return f"图像格式无效或损坏: {str(img_error)}"

            # 将图片数据转换为base64
            image_base64 = base64.b64encode(image_data).decode('utf-8')

            # 使用LangChain的多模态消息格式
            from langchain_core.messages import HumanMessage

            # 使用用户输入的文字提示，如果没有则使用默认提示
            if user_prompt:
                prompt_text = user_prompt
            else:
                prompt_text = "请准确识别这张图片中的所有文字内容，包括中文、英文、数字和标点符号。返回纯文本结果，不要添加任何解释或评论。如果图片中没有文字，请返回'未检测到文字'。"

            # 根据图片格式设置正确的MIME类型
            mime_type = f"image/{image_format}"

            # 创建包含图片的多模态消息
            message = HumanMessage(
                content=[
                    {"type": "text", "text": prompt_text},
                    {"type": "image_url", "image_url": {
                        "url": f"data:{mime_type};base64,{image_base64}"}}
                ]
            )

            # 使用LangChain的统一调用接口
            response = await self.client.ainvoke([message])

            # 检查响应是否有效
            if not response or not hasattr(response, 'content'):
                return "OCR服务暂时不可用，请稍后再试"

            return response.content

        except Exception as e:
            # 提供更友好的错误信息
            error_msg = str(e)
            if "choices" in error_msg and "null" in error_msg:
                return "OCR服务配置有误，请检查API密钥和模型配置"
            elif "401" in error_msg or "unauthorized" in error_msg.lower():
                return "OCR API密钥无效或已过期，请检查配置"
            elif "404" in error_msg or "not found" in error_msg.lower():
                return "OCR API端点或模型不存在，请检查配置"
            elif "timeout" in error_msg.lower():
                return "OCR服务请求超时，请稍后重试"
            elif "image" in error_msg.lower() and "data" in error_msg.lower():
                return "图像数据格式不正确，请重新上传有效的图片文件"
            else:
                return f"OCR服务暂时不可用: {error_msg}"


# 创建LangGraph工具
@tool
def ocr_recognize_tool(image_data: str, user_prompt: str = None) -> str:
    """识别图片中的文字"""
    import asyncio
    ocr_model = GeneralOCR()

    # 检查是否已经有运行的事件循环
    try:
        loop = asyncio.get_running_loop()
        # 如果已经有运行的事件循环，使用create_task

        async def run_ocr():
            return await ocr_model.recognize_text(image_data.encode(), user_prompt=user_prompt)

        # 在当前事件循环中运行
        task = loop.create_task(run_ocr())
        return loop.run_until_complete(task)
    except RuntimeError:
        # 没有运行的事件循环，使用asyncio.run
        return asyncio.run(ocr_model.recognize_text(image_data.encode(), user_prompt=user_prompt))


@tool
def text_chat_tool(message: str) -> str:
    """进行文本对话"""
    import asyncio
    openai_model = OpenAIChatModel()
    messages = [{"role": "user", "content": message}]

    # 检查是否已经有运行的事件循环
    try:
        loop = asyncio.get_running_loop()
        # 如果已经有运行的事件循环，使用create_task

        async def run_chat():
            return await openai_model.chat(messages)

        # 在当前事件循环中运行
        task = loop.create_task(run_chat())
        return loop.run_until_complete(task)
    except RuntimeError:
        # 没有运行的事件循环，使用asyncio.run
        return asyncio.run(openai_model.chat(messages))


class LangGraphMultimodalAgent:
    """基于LangGraph 1.0的多模态聊天代理"""

    def __init__(self):
        self.openai_model = OpenAIChatModel()
        self.ocr_model = GeneralOCR()

        # 创建工具列表
        self.tools = [ocr_recognize_tool, text_chat_tool]

        # 系统提示词，指导AI如何输出表格数据和处理图片消息
        self.system_prompt = """你是一个智能助手，能够处理各种类型的数据，包括文本、图片和文件。请遵循以下指导原则：

            ## 图片处理原则
            1. 当用户上传图片时，系统已经通过OCR技术识别了图片中的文字内容
            2. 你接收到的消息格式为："用户上传了一张图片，图片中的文字是：[OCR识别结果]\n\n用户说：[用户输入]"
            3. 请基于OCR识别结果和用户输入进行智能回复，不要提及"无法查看图片"或"无法分析图片"等限制
            4. 将OCR识别结果作为已知信息，结合用户的问题提供有价值的回答

            ## 表格输出格式
            当用户请求表格数据时，请按照以下格式输出：

            1. 对于简单的数据表格，使用Markdown表格格式：
            ```
            | 列1 | 列2 | 列3 |
            |-----|-----|-----|
            | 数据1 | 数据2 | 数据3 |
            | 数据4 | 数据5 | 数据6 |
            ```

            2. 对于复杂的数据，使用HTML表格格式：
            ```html
            <table>
            <thead>
                <tr>
                <th>列1</th>
                <th>列2</th>
                <th>列3</th>
                </tr>
            </thead>
            <tbody>
                <tr>
                <td>数据1</td>
                <td>数据2</td>
                <td>数据3</td>
                </tr>
                <tr>
                <td>数据4</td>
                <td>数据5</td>
                <td>数据6</td>
                </tr>
            </tbody>
            </table>
            ```

            3. 如果用户明确要求表格形式输出，优先使用HTML表格格式，确保前端能够正确渲染。

            4. 对于不适合表格格式的数据，请使用自然语言描述。

            ## 通用原则
            - 保持回答自然、流畅、有帮助
            - 基于提供的所有信息进行回答，不要提及技术限制
            - 对于图片内容，直接基于OCR识别结果进行回答
            - 如果OCR识别结果不完整或有误，可以合理推断或询问澄清

            请根据用户的具体需求选择合适的输出格式和回答方式。"""

    async def _process_input_node(self, state: AgentState) -> AgentState:
        """处理输入节点"""
        # 根据消息类型处理输入
        if state["message_type"] == "image" and state["file_data"]:
            # 如果是图片，先进行OCR识别，使用用户输入的文字作为提示
            ocr_result = await self.ocr_model.recognize_text(state["file_data"], user_prompt=state["current_input"])
            processed_input = f"用户上传了一张图片，图片中的文字是：{ocr_result}\\n\\n用户说：{state['current_input']}"
            # 保存OCR识别的原始内容
            state["ocr_raw_content"] = ocr_result
        elif state["message_type"] == "file" and state["file_data"]:
            # 处理文件
            if len(state["file_data"]) < 100000:
                try:
                    file_content = state["file_data"].decode('utf-8')
                    processed_input = f"用户上传了一个文件，文件内容是：{file_content}\\n\\n用户说：{state['current_input']}"
                except:
                    processed_input = f"用户上传了一个文件（无法读取文本内容）\\n\\n用户说：{state['current_input']}"
            else:
                processed_input = f"用户上传了一个较大的文件\\n\\n用户说：{state['current_input']}"
            # 文件消息没有OCR内容
            state["ocr_raw_content"] = ""
        else:
            # 纯文本消息
            processed_input = state["current_input"]
            # 文本消息没有OCR内容
            state["ocr_raw_content"] = ""

        # 更新状态，为React代理准备正确的消息格式
        state["messages"] = [{"role": "user", "content": processed_input}]
        return state

    def _generate_response_node(self, state: AgentState) -> AgentState:
        """生成响应节点"""
        # 从代理输出中提取响应
        if "agent" in state and "output" in state["agent"]:
            state["response"] = state["agent"]["output"]
        elif "messages" in state and state["messages"]:
            # 从消息历史中提取最后一个AI消息
            last_message = state["messages"][-1]

            # 处理不同类型的消息对象
            try:
                if hasattr(last_message, 'content'):
                    # 如果是AIMessage对象
                    state["response"] = last_message.content
                elif isinstance(last_message, dict):
                    # 如果是字典格式的消息
                    if last_message.get("role") == "assistant":
                        state["response"] = last_message.get("content", "")
                    else:
                        # 如果不是AI消息，尝试从其他字段获取内容
                        state["response"] = last_message.get(
                            "content", last_message.get("text", ""))
                elif hasattr(last_message, 'get'):
                    # 其他支持get方法的对象
                    if last_message.get("role") == "assistant":
                        state["response"] = last_message.get("content", "")
                    else:
                        state["response"] = last_message.get(
                            "content", last_message.get("text", ""))
                else:
                    # 如果无法识别消息格式，使用默认响应
                    state["response"] = "抱歉，未能生成响应"
            except Exception as e:
                # 捕获消息格式处理错误
                state["response"] = "抱歉，处理消息时出现格式错误"

        # 如果仍未设置响应，使用默认值
        if not state.get("response"):
            state["response"] = "抱歉，未能生成响应"

        return state

    async def process_message_stream_with_context(self, message_history: List[Dict[str, str]],
                                                  message_type: str = "text", file_data: bytes = None,
                                                  enable_search: bool = True):
        """流式处理用户消息（带上下文历史），支持联网搜索"""
        try:
            # 提取当前用户消息
            current_message = message_history[-1]["content"] if message_history else ""

            # 检查用户消息是否包含表格请求关键词
            table_keywords = ["表格", "table", "列表", "清单", "数据表", "以表格形式"]
            requires_table = any(keyword in current_message.lower()
                                 for keyword in table_keywords)

            # 创建初始状态
            initial_state = AgentState(
                messages=[],
                current_input=current_message,
                message_type=message_type,
                file_data=file_data or b"",
                response="",
                ocr_raw_content=""
            )

            # 处理输入（OCR等）- 使用异步版本
            processed_state = await self._process_input_node(initial_state)

            # 如果是图片消息且有OCR识别内容，先输出OCR原始内容
            if message_type == "image" and file_data and processed_state.get("ocr_raw_content"):
                ocr_content = processed_state["ocr_raw_content"]
                if ocr_content and ocr_content.strip():
                    # 输出OCR识别的原始内容
                    yield f"=== OCR识别原始内容 ===\n{ocr_content}\n=== OCR识别结束 ===\n\n"

            # 构建包含历史消息的完整上下文
            context_messages = []

            # 添加系统提示词（如果需要表格）
            if requires_table:
                context_messages.append(
                    {"role": "system", "content": self.system_prompt})

            for msg in message_history[:-1]:  # 排除当前消息
                context_messages.append(msg)

            # 添加处理后的当前消息
            context_messages.append({
                "role": "user",
                "content": processed_state["messages"][0]["content"] if processed_state["messages"] else current_message
            })

            # 使用流式聊天，支持搜索开关
            async for chunk in self.openai_model.chat_stream(context_messages, enable_search):
                yield chunk

        except Exception as e:
            # 捕获具体的错误信息
            error_msg = str(e)
            if "Message field is required" in error_msg or "20014" in error_msg:
                yield "AI服务配置错误：API端点期望的消息格式不兼容。请检查API配置或使用兼容的API服务。"
            else:
                yield f"抱歉，处理消息时出现错误：{error_msg}"


# 全局AI代理实例 - 使用LangChain 1.0版本
ai_agent = LangGraphMultimodalAgent()
